import torch
import os
from typing import Dict, Union


def get_device_info(device: torch.device) -> Dict[str, Union[str, int]]:
    """获取详细的设备信息

    Args:
        device: torch.device对象

    Returns:
        包含设备详细信息的字典
    """
    info = {
        'device_type': device.type,
        'device_index': device.index if device.index is not None else 0
    }

    if device.type == 'cuda':
        info.update({
            'gpu_name': torch.cuda.get_device_name(device),
            'cuda_version': torch.version.cuda,
            'total_memory': f"{torch.cuda.get_device_properties(device).total_memory / 1024 ** 3:.2f}GB"
        })
    return info


def log_device_info(logger):
    """记录设备信息到日志"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    info = get_device_info(device)

    logger.info("=" * 50)
    logger.info("Device Configuration:")
    for k, v in info.items():
        logger.info(f"{k.upper():<15}: {v}")
    logger.info("=" * 50)